"""
广义稀疏点积-对角矩阵乘法操作

提供与DGL兼容的gsddmm操作，使用PyG作为后端
"""

import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import sys
import os

# 添加路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(current_dir))
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

from ops.src.graph.graph import is_dgl_graph


def gsddmm(graph, op, lhs_data, rhs_data, lhs_target='u', rhs_target='v'):
    # print("===================gsddmm===================")
    """
    广义稀疏点积-对角矩阵乘法
    
    参数:
        graph: Graph对象或DGL图
        op: 操作类型，如 'mul', 'add', 'sub', 'div', 'dot'
        lhs_data: 左侧数据，通常是节点特征
        rhs_data: 右侧数据，通常是节点特征
        lhs_target: 左侧数据目标，'u' 表示源节点，'v' 表示目标节点
        rhs_target: 右侧数据目标，'u' 表示源节点，'v' 表示目标节点
    
    返回:
        计算结果，边特征
    """
    # 检查图是否为DGL图
    if is_dgl_graph(graph):
        # DGL图
        import dgl.ops as dgl_ops
        return dgl_ops.gsddmm(graph, op, lhs_data, rhs_data, lhs_target, rhs_target)
    else:
        # PyG风格的Graph
        src, dst = graph.edge_index
        
        # 根据目标选择节点
        if lhs_target == 'u':
            lhs_idx = src
        elif lhs_target == 'v':
            lhs_idx = dst
        else:
            raise ValueError(f"Unsupported lhs_target: {lhs_target}")
        
        if rhs_target == 'u':
            rhs_idx = src
        elif rhs_target == 'v':
            rhs_idx = dst
        else:
            raise ValueError(f"Unsupported rhs_target: {rhs_target}")
        
        # 获取对应节点的特征
        lhs_features = lhs_data[lhs_idx]
        rhs_features = rhs_data[rhs_idx]
        
        # 执行操作
        if op == 'mul':
            return lhs_features * rhs_features
        elif op == 'add':
            return lhs_features + rhs_features
        elif op == 'sub':
            return lhs_features - rhs_features
        elif op == 'div':
            return lhs_features / (rhs_features + 1e-6)  # 添加小值防止除零
        elif op == 'dot':
            # 点积操作，要求特征维度匹配
            return torch.sum(lhs_features * rhs_features, dim=-1, keepdim=True)
        else:
            raise ValueError(f"Unsupported operation: {op}") 